#include "ekf.h"

Ekf::Ekf():
	P_(3,3) {
	P_ << 0, 0, 0,
			  0, 0, 0,
				0, 0, 0;
}

void Ekf::AssociateLandmarks(const robot_control::LandmarksConstPtr & msg,
		double pose_x, double pose_y, double angle) {
	std::vector<Landmark>::iterator it;
	new_size_ = msg->new_size;

	ROS_INFO("Associate %d new, %d old", new_size_,
			msg->id.size() - new_size_);

	visible_landmarks_.clear();
	bearings_old_.clear();
	bearings_old_id_.clear();
	bearings_new_.clear();

	x_ = pose_x;
	y_ = pose_y;
	angle_ = angle;

	for (size_t i = 0; i < msg->id.size(); i++) {
			Landmark lm;
			lm.feature_id = msg->id[i];
			lm.x = msg->x[i];
			lm.y = msg->y[i];
			lm.z = msg->z[i];
			visible_landmarks_.push_back(lm);

			double dx = lm.x - x_;
			double dy = lm.y - y_;
			double phi = atan(dy/dx) - angle_;
			Coordinates bearing(sqrt(pow(dx,2) + pow(dy,2)), phi);
			if (i < new_size_) {
				bearings_new_.push_back(bearing);
			} else {
				bearings_old_.push_back(bearing);
				bearings_old_id_.push_back(lm.feature_id);
			}
	}

}

void Ekf::Predict() {

	double delta_angle = rotation_speed_ * time_elapsed_;
	double s = sin(angle_ + delta_angle);
	double c = cos(angle_ + delta_angle);
	double vts = velocity_ * time_elapsed_ * s;
	double vtc = velocity_ * time_elapsed_ * c;

	Eigen::Matrix3d Gv;
	Eigen::MatrixXd Gu(2, 3);
	Gv << 1, 0, -vts,
		    0, 1, vtc,
				0, 0, 1;
	Gu << time_elapsed_ * c, -vts,
		    time_elapsed_ * s, vtc,
				time_elapsed_ * sin(angle_)/WB, velocity_*time_elapsed_*cos(angle_)/WB;

	Eigen::Matrix3d P3;
	Eigen::MatrixXd Pright(3, P_.cols()-3);
	
	for (int i = 0; i < 3; i++) {
		for (int j = 0; j < 3; j++) {
			P3(i, j) = P_(i, j);
		}
	}
	for (int i = 0, x = 0; i < 3; i++, x++) {
		for (int j = 3, y = 0; j < P_.cols(); j++, y++) {
      Pright(x, y) = P_(i, j);
		}
	}
	
	P3 = Gv * P3 * Gv.transpose();
	
	//JJ The following need just when we have features, j>4, do we need to add this if(size(P,1)>3)?
	// N, if j<4 the for will not be used so it is safe
	Pright = Gv * Pright;

	for (int i = 0, x = 0; i < 3; i++, x++) {
		for (int j = 3,  y = 0; j < P_.cols(); j++, y++) {
      P_(i,j) = Pright(x, y); // DO we need this line? because we already compute P(1:3,4:end), we just need to fill in P(4:end,1:3)
			// N we computed it but we did not put it inside P_
      P_(j,i) = Pright.transpose()(y, x);
		}
	}

	x_ = x_ + vtc;
	y_ = y_ + vts;
  angle_ = pi_to_pi(angle_ + velocity_ * time_elapsed_ * sin(angle_) / WB);

}

void Ekf::Update() {

  for (size_t i = 0; i < bearings_old_.size(); i++) { //i is the number of associated landmarks, as the Zf_bearing in the Matlab
  
  //-----------fuction observe_model--------------
		printf("%lf,%f - %lf,%lf\n", visible_landmarks_[i+new_size_].x,
				visible_landmarks_[i+new_size_].y, x_, y_);
    double dx = visible_landmarks_[i+new_size_].x - x_; //i think this x and y should be the predicted_x and predicted_y , 
                                                 //but this predicted_x,y are updated after each landmark
    
    double dy = visible_landmarks_[i+new_size_].y - y_;
		double d2 = pow(dx, 2) + pow(dy, 2);
	  double d = sqrt(d2);
		double xd = dx/d;
		double yd = dy/d;
		double xd2 = dx/d2;
		double yd2 = dy/d2;

		printf("<<<<%lf,%lf\n", dx, dy);
		
		Eigen::MatrixXd z(2, 1); 
		z << d, atan(dy/dx) - angle_;
		
		Eigen::MatrixXd H = Eigen::MatrixXd::Zero(2, P_.cols());
		H(0,0) = -xd; H(0,1) = -yd; H(0,2) = 0;
		H(1,0) = yd2; H(1,1) = -xd2; H(1,2) = -1;
		// The landmark id will always start from 1
    H(0,bearings_old_id_[i]+2) = xd;
		H(0, bearings_old_id_[i]+3) = yd;
    H(1,bearings_old_id_[i]+2) = -yd2;
		H(1, bearings_old_id_[i]+3) = xd2;
//---------- fuction finish ----------------


		Eigen::MatrixXd v(2,1);
    v << bearings_old_[i].first - z(0,0),
			   pi_to_pi(bearings_old_[i].second - z(1,0));
			   
//--------- fuction cholesky------			   

		Eigen::MatrixXd Pht(P_.cols(), 2);
		std::cout << "P: " << P_ <<std::endl;
		std::cout << "H: " << H <<std::endl;
		Pht = P_ * H.transpose();
		Eigen::Matrix2d S;
		S = H*Pht; 
		S = (S+S.transpose()) * 0.5;

		Eigen::Matrix2d SChol;
		//TODO check this mysterious function
		//S.ldlt().solve(S, &SChol);

		Eigen::MatrixXd W(P_.cols(), 2);
		Eigen::MatrixXd W1(P_.cols(), 2);
    W1 = Pht * SChol.inverse();

		Eigen::Matrix2d SCholInv;
		SCholInv = SChol.inverse();
		W = W1 * SCholInv.transpose();

    Eigen::MatrixXd x_up = Eigen::MatrixXd::Zero(P_.cols(), 1);
		x_up = W*v; 
		std::cout << "S: " << S <<std::endl;
		std::cout << "Schol: " << SChol <<std::endl;
		x_ += x_up(0,0); 
		y_ += x_up(1,0);
		angle_ += x_up(2,0);
		for (int i = 3; i < x_up.rows(); i+=2) {
			visible_landmarks_[i].x += x_up(i, 0);
			visible_landmarks_[i].y += x_up(i+1, 0);
		}
		P_ = P_ - W1*W1.transpose();
	}
}



//For augument, we just use the matched features not associated part for adding to the landmark list and compute the P_
//Because this P_ related to the SLAM state (robot pose + all the landmarks)
void Ekf::Augment() { // this is the output of the updated robot pose
	for (size_t ii = 0; ii < bearings_new_.size(); ii++) { //this loop will just use the new landmark, like the Zn_bearing in Matlab
		
		Eigen::MatrixXd Pnew(P_.cols()+2,P_.cols()+2);
		double s = sin(angle_+bearings_new_[ii].second);
		double c = cos(angle_+bearings_new_[ii].second);
		double r = bearings_new_[ii].first;
		printf("s:%lf c:%lf r:%lf\n", s, c, r);
		
		Eigen::MatrixXd Gv(2,3);
		Eigen::Matrix2d Gz;
		Gv << 1, 0, -r*s,
		   	  0, 1, r*c;
		Gz << c, -r*s,
			    s, r*c;

		for (int x = 0; x < P_.cols(); x++) {
			for (int y = 0; y < P_.cols(); y++) {
				Pnew(x, y) = P_(x, y);
			}
		}

		Eigen::Matrix3d P3;
		Eigen::Matrix2d P2;
		Eigen::MatrixXd Pdown(2, 3);
		for (int i = 0; i < 3; i++) {
			for (int j = 0; j < 3; j++) {
				P3(i, j) = P_(i, j);
			}
		}
		P2 = Gv * P3 * Gv.transpose();
		
		for (int i = 0, x = P_.cols(); i < 2; i++, x++) {
			for (int j = 0, y = P_.cols(); j < 2; j++, y++) {
				Pnew(x, y) = P2(i, j);
			}
		}
		
		Pdown = Gv * P3;
		Eigen::MatrixXd Pdown_tr = Pdown.transpose();

		for (int i = 0, x = P_.cols(); i < 2; i++, x++) {
			for (int j = 0, y = 0; j < 3; j++, y++) {
				Pnew(x,y) = Pdown(i, j);
				Pnew(y,x) = Pdown_tr(j, i);
			}
		}

		if (P_.cols() > 3) {
			Eigen::MatrixXd Pother(3, Pnew.cols());
			Eigen::MatrixXd Pxcorr(2, Pnew.cols());
			for (int i = 0, x = 0; x < 3; i++, x++) {
				for (int j = 0, y = 3; y < Pnew.cols(); j++, y++) {
					Pother(i,j) = Pnew(x,y);
				}
			}
			Pxcorr = Gv * Pother;
			Eigen::MatrixXd Pxcorr_tr = Pxcorr.transpose();

			for (int i = 0, x = P_.cols(); i < 2; i++, x++) { 
				for (int j = 0, y = 3; y < Pxcorr.cols(); y++, j++) {
					Pnew(x, y) = Pxcorr(i, j);
					Pnew(y, x) = Pxcorr_tr(j, i);
				}
			}
		}

		P_ = Pnew;
		std::cout << "round: "<< ii << std::endl;
		std::cout << P_ << std::endl;
	}
}

void Ekf::CorrectPose(double *x, double *y, double *angle,
		double velocity, double rotation_speed, double time_elapsed) {
	x_ = *x;
	y_ = *y;
	angle_ = *angle;
	velocity_ = velocity;
	rotation_speed_ = rotation_speed;
	time_elapsed_ = time_elapsed;

	printf(">>>>%lf, %lf, %lf, %lf, %lf, %lf\n", x_, y_, angle_, velocity_,
			rotation_speed_, time_elapsed);

	Predict();
	Update();
	Augment();

	*x = x_;
	*y = y_;
	*angle = angle_;
}

double Ekf::pi_to_pi(const double angle) const {
	double new_angle = fmod(angle, 2* M_PI);
	if (new_angle > M_PI)
		return new_angle - 2 * M_PI;
	if (new_angle < -M_PI)
		return new_angle + 2 * M_PI;

	return new_angle;
}
